package gbase

import (
	"sync"
	"time"
)

// 通用令牌桶限流器
// Limiter 是通用限流器接口
// 可扩展支持多种限流算法
// 目前实现令牌桶算法

type Limiter interface {
	// TryAcquire 尝试获取一个令牌，立即返回是否成功
	TryAcquire() bool
	// WaitAcquire 阻塞直到获取到一个令牌
	WaitAcquire()
	// SetRate 动态调整速率
	SetRate(rate float64)
	// SetCapacity 动态调整桶容量
	SetCapacity(capacity int)
}

type TokenBucket struct {
	capacity int       // 桶容量
	rate     float64   // 每秒生成令牌数
	tokens   float64   // 当前令牌数
	last     time.Time // 上次补充令牌时间
	mu       sync.Mutex
}

// NewTokenBucket 创建一个令牌桶限流器
func NewTokenBucket(rate float64, capacity int) *TokenBucket {
	if capacity <= 0 {
		capacity = 1
	}
	if rate <= 0 {
		rate = 1
	}
	return &TokenBucket{
		capacity: capacity,
		rate:     rate,
		tokens:   float64(capacity),
		last:     time.Now(),
	}
}

// TryAcquire 尝试获取一个令牌，立即返回是否成功
func (tb *TokenBucket) TryAcquire() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	tb.refill()
	if tb.tokens >= 1 {
		tb.tokens -= 1
		return true
	}
	return false
}

// WaitAcquire 阻塞直到获取到一个令牌
func (tb *TokenBucket) WaitAcquire() {
	for {
		if tb.TryAcquire() {
			return
		}
		time.Sleep(time.Millisecond * 10)
	}
}

// SetRate 动态调整速率
func (tb *TokenBucket) SetRate(rate float64) {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	tb.refill()
	if rate > 0 {
		tb.rate = rate
	}
}

// SetCapacity 动态调整桶容量
func (tb *TokenBucket) SetCapacity(capacity int) {
	tb.mu.Lock()
	defer tb.mu.Unlock()
	tb.refill()
	if capacity > 0 {
		tb.capacity = capacity
		if tb.tokens > float64(capacity) {
			tb.tokens = float64(capacity)
		}
	}
}

// 补充令牌
func (tb *TokenBucket) refill() {
	now := time.Now()
	delta := now.Sub(tb.last).Seconds()
	if delta > 0 {
		add := delta * tb.rate
		tb.tokens += add
		if tb.tokens > float64(tb.capacity) {
			tb.tokens = float64(tb.capacity)
		}
		tb.last = now
	}
}

// 确保接口实现
var _ Limiter = (*TokenBucket)(nil)
